from flask import Flask, render_template, request, redirect, url_for
from werkzeug.utils import secure_filename
import os
from predict import load_model, predict_image
import matplotlib.pyplot as plt
import io
import base64
import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端
app = Flask(__name__)

# 配置
UPLOAD_FOLDER = 'static/uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 确保上传文件夹存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

# 加载模型
model = load_model('garbage_classification.pth')

# 检查文件扩展名
def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

# 生成概率图
def plot_to_img(fig):
    buf = io.BytesIO()
    fig.savefig(buf, format='png', bbox_inches='tight')
    buf.seek(0)
    img_data = base64.b64encode(buf.read()).decode('utf-8')
    plt.close(fig)
    return img_data

@app.route('/', methods=['GET', 'POST'])
def upload_file():
    if request.method == 'POST':
        # 检查是否有文件上传
        if 'file' not in request.files:
            return redirect(request.url)
        file = request.files['file']
        
        if file.filename == '':
            return redirect(request.url)
        
        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)
            
            # 进行预测
            predicted_class, confidence, all_probs = predict_image(filepath, model)
            
            # 生成概率图
            fig, ax = plt.subplots()
            ax.bar(all_probs.keys(), all_probs.values())
            ax.set_title('Classification Probabilities')
            ax.set_ylabel('Probability')
            plt.xticks(rotation=45)
            plot_url = plot_to_img(fig)
            
            return render_template('index.html', 
                                filename=filename,
                                predicted_class=predicted_class,
                                confidence=confidence*100,
                                plot_url=plot_url)
    
    return render_template('index.html')

if __name__ == '__main__':
    app.run(debug=True)